import os
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.vectorstores import Chroma
from ics.backend.data_process_answer.text_split import text_splitter

model_embedding = "G:/M3E_model/xrunda/m3e-base"
model_kwargs = {'device': 'cuda'}

hf = HuggingFaceBgeEmbeddings(
    encode_kwargs={'normalize_embeddings': True},
    model_name=model_embedding,
    model_kwargs=model_kwargs,
    query_instruction="为这个句子生成表示以用于检索相关文章："
)

def create_or_update_vectordb(documents, persist_directory):
    """创建向量数据库或加载已有的数据库，并在需要时添加新的数据。"""
    if not os.path.exists(persist_directory):
        # 如果数据库不存在，则创建新的数据库并持久化
        db = Chroma.from_documents(
            documents=documents,
            embedding=hf,
            ids=None,  # 自动生成 ids
            collection_metadata={"hnsw:space": "cosine"},
            persist_directory=persist_directory,
        )
    else:
        # 如果数据库已存在，则加载数据库
        db = Chroma.load(persist_directory=persist_directory, embedding_function=hf)
        # 获取现有数据库中的所有文档 ID
        existing_ids = set(db.get()['ids'])
        # 假设每个文档都有一个唯一的 id
        new_documents = [doc for doc in documents if doc.metadata['id'] not in existing_ids]
        if new_documents:
            # 将新的文档添加到数据库中
            db.add_documents(new_documents)
    return db
if __name__ == "__main__":
    # 加载并分割文本
    documents = text_splitter("D:/Anything/file/s1.md")

    # 为每个文档生成一个唯一的 id
    for i, doc in enumerate(documents):
        doc.metadata = {'id': f'doc_{i}'}

    # 创建或更新向量数据库
    db = create_or_update_vectordb(documents, './db/chroma_db')

    print("向量数据库更新完成。")
